{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4720f86f",
   "metadata": {},
   "source": [
    "# Tutorial 04: computation of the inf-sup constant for a Stokes problem discretization\n",
    "\n",
    "In this tutorial we compare the computation of the inf-sup constant of a Stokes problem using standard assembly with mixed function spaces and block assembly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d95ef51",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4abe105",
   "metadata": {},
   "outputs": [],
   "source": [
    "import basix.ufl\n",
    "import dolfinx.cpp\n",
    "import dolfinx.fem\n",
    "import dolfinx.mesh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import numpy.typing\n",
    "import petsc4py.PETSc\n",
    "import slepc4py.SLEPc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f011a69a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "085c9482",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffd40bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "342c905d-4564-4e89-ad17-199121536bd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b30f65b8-8aa7-46f3-bef8-5faffd663604",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wall(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Determine the position of the wall.\"\"\"\n",
    "    return np.logical_or(  # type: ignore[no-any-return]\n",
    "        x[1] < 0 + np.finfo(float).eps, x[1] > 1 - np.finfo(float).eps)\n",
    "\n",
    "\n",
    "boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim - 1, wall)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c771f69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7c89f1e",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "694b5737",
   "metadata": {},
   "outputs": [],
   "source": [
    "V_element = basix.ufl.element(\"Lagrange\", mesh.basix_cell(), 2, shape=(mesh.geometry.dim, ))\n",
    "Q_element = basix.ufl.element(\"Lagrange\", mesh.basix_cell(), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c69d76cf",
   "metadata": {},
   "source": [
    "### Auxiliary function for eigenvector normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b449fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(u1: dolfinx.fem.Function, u2: dolfinx.fem.Function, p: dolfinx.fem.Function) -> None:\n",
    "    \"\"\"Normalize an eigenvector.\"\"\"\n",
    "    scaling_operations: list[tuple[  # type: ignore[no-any-unimported]\n",
    "        dolfinx.fem.Function, typing.Callable[[dolfinx.fem.Function], ufl.Form],\n",
    "        typing.Callable[[petsc4py.PETSc.ScalarType], petsc4py.PETSc.ScalarType]\n",
    "    ]] = [\n",
    "        # Scale functions with a W^{1,1} (for velocity) or L^1 (for pressure) norm to take away\n",
    "        # possible sign differences.\n",
    "        (u1, lambda u: (u.dx(0) + u.dx(1)) * ufl.dx, lambda x: x),\n",
    "        (u2, lambda u: (u.dx(0) + u.dx(1)) * ufl.dx, lambda x: x),\n",
    "        (p, lambda p: p * ufl.dx, lambda x: x),\n",
    "        # Normalize functions with a H^1 (for velocity) or L^2 (for pressure) norm.\n",
    "        (u1, lambda u: ufl.inner(ufl.grad(u), ufl.grad(u)) * ufl.dx, lambda x: np.sqrt(x)),\n",
    "        (u2, lambda u: ufl.inner(ufl.grad(u), ufl.grad(u)) * ufl.dx, lambda x: np.sqrt(x)),\n",
    "        (p, lambda p: ufl.inner(p, p) * ufl.dx, lambda x: np.sqrt(x))\n",
    "    ]\n",
    "    for (function, bilinear_form, postprocess) in scaling_operations:\n",
    "        scalar = postprocess(mesh.comm.allreduce(\n",
    "            dolfinx.fem.assemble_scalar(dolfinx.fem.form(bilinear_form(function))), op=mpi4py.MPI.SUM))\n",
    "        function.x.petsc_vec.scale(1. / scalar)\n",
    "        function.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "673c932c",
   "metadata": {},
   "source": [
    "### Standard formulation using a mixed function space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2900acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_monolithic() -> tuple[np.float64, dolfinx.fem.Function, dolfinx.fem.Function, dolfinx.fem.Function]:\n",
    "    \"\"\"Run standard formulation using a mixed function space.\"\"\"\n",
    "    # Function spaces\n",
    "    W_element = basix.ufl.mixed_element([V_element, Q_element])\n",
    "    W = dolfinx.fem.functionspace(mesh, W_element)\n",
    "\n",
    "    # Test and trial functions: monolithic\n",
    "    vq = ufl.TestFunction(W)\n",
    "    (v, q) = ufl.split(vq)\n",
    "    up = ufl.TrialFunction(W)\n",
    "    (u, p) = ufl.split(up)\n",
    "\n",
    "    # Variational forms\n",
    "    lhs = (ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx - ufl.inner(p, ufl.div(v)) * ufl.dx\n",
    "           - ufl.inner(ufl.div(u), q) * ufl.dx)\n",
    "    rhs = - ufl.inner(p, q) * ufl.dx\n",
    "\n",
    "    # Define restriction for DOFs associated to homogenous Dirichlet boundary conditions\n",
    "    dofs_W = np.arange(0, W.dofmap.index_map.size_local + W.dofmap.index_map.num_ghosts)\n",
    "    W0 = W.sub(0)\n",
    "    V, _ = W0.collapse()\n",
    "    bdofs_V = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundary_facets)[0]\n",
    "    restriction = multiphenicsx.fem.DofMapRestriction(W.dofmap, np.setdiff1d(dofs_W, bdofs_V))\n",
    "\n",
    "    # Assemble lhs and rhs matrices\n",
    "    A = multiphenicsx.fem.petsc.assemble_matrix(\n",
    "        dolfinx.fem.form(lhs), restriction=(restriction, restriction))\n",
    "    A.assemble()\n",
    "    B = multiphenicsx.fem.petsc.assemble_matrix(\n",
    "        dolfinx.fem.form(rhs), restriction=(restriction, restriction))\n",
    "    B.assemble()\n",
    "\n",
    "    # Solve\n",
    "    eps = slepc4py.SLEPc.EPS().create(mesh.comm)\n",
    "    eps.setOperators(A, B)\n",
    "    eps.setProblemType(slepc4py.SLEPc.EPS.ProblemType.GNHEP)\n",
    "    eps.setDimensions(1, petsc4py.PETSc.DECIDE, petsc4py.PETSc.DECIDE)\n",
    "    eps.setWhichEigenpairs(slepc4py.SLEPc.EPS.Which.TARGET_REAL)\n",
    "    eps.setTarget(1.e-5)\n",
    "    eps.getST().setType(slepc4py.SLEPc.ST.Type.SINVERT)\n",
    "    eps.getST().getKSP().setType(\"preonly\")\n",
    "    eps.getST().getKSP().getPC().setType(\"lu\")\n",
    "    eps.getST().getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    eps.solve()\n",
    "    assert eps.getConverged() >= 1\n",
    "\n",
    "    # Extract leading eigenvalue and eigenvector\n",
    "    vr = dolfinx.cpp.fem.petsc.create_vector_block([(restriction.index_map, restriction.index_map_bs)])\n",
    "    vi = dolfinx.cpp.fem.petsc.create_vector_block([(restriction.index_map, restriction.index_map_bs)])\n",
    "    eigv = eps.getEigenpair(0, vr, vi)\n",
    "    r, i = eigv.real, eigv.imag\n",
    "    assert abs(i) < 1.e-10\n",
    "    assert r > 0., \"r = \" + str(r) + \" is not positive\"\n",
    "    print(\"Inf-sup constant (monolithic): \", np.sqrt(r))\n",
    "\n",
    "    # Transform eigenvector into eigenfunction\n",
    "    r_fun = dolfinx.fem.Function(W)\n",
    "    vr.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "    with r_fun.x.petsc_vec.localForm() as r_fun_local, \\\n",
    "            multiphenicsx.fem.petsc.VecSubVectorWrapper(vr, W.dofmap, restriction) as vr_wrapper:\n",
    "        r_fun_local[:] = vr_wrapper\n",
    "    u_fun = r_fun.sub(0).collapse()\n",
    "    (u_fun_1, u_fun_2) = (u_fun.sub(0).collapse(), u_fun.sub(1).collapse())\n",
    "    p_fun = r_fun.sub(1).collapse()\n",
    "    normalize(u_fun_1, u_fun_2, p_fun)\n",
    "\n",
    "    eps.destroy()\n",
    "    return (r, u_fun_1, u_fun_2, p_fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4764b572",
   "metadata": {},
   "outputs": [],
   "source": [
    "(eig_m, u_fun_1_m, u_fun_2_m, p_fun_m) = run_monolithic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc0ceb4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_fun_1_m, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cc7bb35",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_fun_2_m, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3336bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p_fun_m, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4aa919b",
   "metadata": {},
   "source": [
    "### Block formulation using two independent function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dddb8966",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_block() -> tuple[np.float64, dolfinx.fem.Function, dolfinx.fem.Function, dolfinx.fem.Function]:\n",
    "    \"\"\"Run block formulation using two independent function spaces.\"\"\"\n",
    "    # Function spaces\n",
    "    V = dolfinx.fem.functionspace(mesh, V_element)\n",
    "    Q = dolfinx.fem.functionspace(mesh, Q_element)\n",
    "\n",
    "    # Test and trial functions\n",
    "    (v, q) = (ufl.TestFunction(V), ufl.TestFunction(Q))\n",
    "    (u, p) = (ufl.TrialFunction(V), ufl.TrialFunction(Q))\n",
    "\n",
    "    # Variational forms\n",
    "    lhs = [[ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx, - ufl.inner(p, ufl.div(v)) * ufl.dx],\n",
    "           [- ufl.inner(ufl.div(u), q) * ufl.dx, None]]\n",
    "    rhs = [[None, None],\n",
    "           [None, - ufl.inner(p, q) * ufl.dx]]\n",
    "    rhs[0][0] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(u, v) * ufl.dx\n",
    "\n",
    "    # Define restriction for DOFs associated to homogenous Dirichlet boundary conditions\n",
    "    dofs_V = np.arange(0, V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts)\n",
    "    bdofs_V = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)\n",
    "    dofs_Q = np.arange(0, Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts)\n",
    "    restriction_V = multiphenicsx.fem.DofMapRestriction(V.dofmap, np.setdiff1d(dofs_V, bdofs_V))\n",
    "    restriction_Q = multiphenicsx.fem.DofMapRestriction(Q.dofmap, dofs_Q)\n",
    "    restriction = [restriction_V, restriction_Q]\n",
    "\n",
    "    # Assemble lhs and rhs matrices\n",
    "    A = multiphenicsx.fem.petsc.assemble_matrix_block(\n",
    "        dolfinx.fem.form(lhs), bcs=[], restriction=(restriction, restriction))\n",
    "    A.assemble()\n",
    "    B = multiphenicsx.fem.petsc.assemble_matrix_block(\n",
    "        dolfinx.fem.form(rhs), bcs=[], restriction=(restriction, restriction))\n",
    "    B.assemble()\n",
    "\n",
    "    # Solve\n",
    "    eps = slepc4py.SLEPc.EPS().create(mesh.comm)\n",
    "    eps.setOperators(A, B)\n",
    "    eps.setProblemType(slepc4py.SLEPc.EPS.ProblemType.GNHEP)\n",
    "    eps.setDimensions(1, petsc4py.PETSc.DECIDE, petsc4py.PETSc.DECIDE)\n",
    "    eps.setWhichEigenpairs(slepc4py.SLEPc.EPS.Which.TARGET_REAL)\n",
    "    eps.setTarget(1.e-5)\n",
    "    eps.getST().setType(slepc4py.SLEPc.ST.Type.SINVERT)\n",
    "    eps.getST().getKSP().setType(\"preonly\")\n",
    "    eps.getST().getKSP().getPC().setType(\"lu\")\n",
    "    eps.getST().getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    eps.solve()\n",
    "    assert eps.getConverged() >= 1\n",
    "\n",
    "    # Extract leading eigenvalue and eigenvector\n",
    "    vr = dolfinx.cpp.fem.petsc.create_vector_block(\n",
    "        [(restriction_.index_map, restriction_.index_map_bs) for restriction_ in restriction])\n",
    "    vi = dolfinx.cpp.fem.petsc.create_vector_block(\n",
    "        [(restriction_.index_map, restriction_.index_map_bs) for restriction_ in restriction])\n",
    "    eigv = eps.getEigenpair(0, vr, vi)\n",
    "    r, i = eigv.real, eigv.imag\n",
    "    assert abs(i) < 1.e-10\n",
    "    assert r > 0., \"r = \" + str(r) + \" is not positive\"\n",
    "    print(\"Inf-sup constant (block): \", np.sqrt(r))\n",
    "\n",
    "    # Transform eigenvector into eigenfunction\n",
    "    (u_fun, p_fun) = (dolfinx.fem.Function(V), dolfinx.fem.Function(Q))\n",
    "    vr.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "    with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(vr, [V.dofmap, Q.dofmap], restriction) as vr_wrapper:\n",
    "        for vr_wrapper_local, component in zip(vr_wrapper, (u_fun, p_fun)):\n",
    "            with component.x.petsc_vec.localForm() as component_local:\n",
    "                component_local[:] = vr_wrapper_local\n",
    "    (u_fun_1, u_fun_2) = (u_fun.sub(0).collapse(), u_fun.sub(1).collapse())\n",
    "    normalize(u_fun_1, u_fun_2, p_fun)\n",
    "\n",
    "    eps.destroy()\n",
    "    return (r, u_fun_1, u_fun_2, p_fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccf10425",
   "metadata": {},
   "outputs": [],
   "source": [
    "(eig_b, u_fun_1_b, u_fun_2_b, p_fun_b) = run_block()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f6fa68d",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_fun_1_b, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21450fc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_fun_2_b, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1d0015",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p_fun_b, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b33084a0",
   "metadata": {},
   "source": [
    "### Error computation between standard and block formulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "439a839c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_error(\n",
    "    eig_m: np.float64, eig_b: np.float64, u_fun_1_m: dolfinx.fem.Function, u_fun_1_b: dolfinx.fem.Function,\n",
    "    u_fun_2_m: dolfinx.fem.Function, u_fun_2_b: dolfinx.fem.Function, p_fun_m: dolfinx.fem.Function,\n",
    "    p_fun_b: dolfinx.fem.Function\n",
    ") -> None:\n",
    "    \"\"\"Compute errors between the mixed and block cases.\"\"\"\n",
    "    err_inf_sup = np.abs(np.sqrt(eig_b) - np.sqrt(eig_m)) / np.sqrt(eig_m)\n",
    "    print(\"Relative error for inf-sup constant equal to\", err_inf_sup)\n",
    "    assert np.isclose(err_inf_sup, 0., atol=1.e-8)\n",
    "    eigenvector_operations: list[tuple[  # type: ignore[no-any-unimported]\n",
    "        dolfinx.fem.Function, dolfinx.fem.Function, typing.Callable[[dolfinx.fem.Function], ufl.Form], str\n",
    "    ]] = [\n",
    "        (u_fun_1_b, u_fun_1_m, lambda u: ufl.inner(ufl.grad(u), ufl.grad(u)) * ufl.dx, \"velocity 1\"),\n",
    "        (u_fun_2_b, u_fun_2_m, lambda u: ufl.inner(ufl.grad(u), ufl.grad(u)) * ufl.dx, \"velocity 2\"),\n",
    "        (p_fun_b, p_fun_m, lambda p: ufl.inner(p, p) * ufl.dx, \"pressure\")\n",
    "    ]\n",
    "    for (fun_b, fun_m, squared_norm_form, component_name) in eigenvector_operations:\n",
    "        err_fun = np.sqrt(mesh.comm.allreduce(\n",
    "            dolfinx.fem.assemble_scalar(dolfinx.fem.form(squared_norm_form(fun_b - fun_m))), op=mpi4py.MPI.SUM))\n",
    "        print(\"Relative error for \", component_name, \"component of eigenvector equal to\", err_fun)\n",
    "        assert np.isclose(err_fun, 0., atol=1.e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "783cb0fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_error(eig_m, eig_b, u_fun_1_m, u_fun_1_b, u_fun_2_m, u_fun_2_b, p_fun_m, p_fun_b)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
